from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
from docx import Document
from docx.shared import Pt
from markdown import markdown
import io
import tempfile
import os
from bs4 import BeautifulSoup
import re

app = FastAPI()

def convert_math_expressions(html_content):
    """转换数学表达式，保持格式"""
    # 处理行内数学公式 \( ... \)
    html_content = re.sub(r'\\\((.*?)\\\)', r'$\1$', html_content)
    # 处理块级数学公式 \[ ... \]
    html_content = re.sub(r'\\\[(.*?)\\\]', r'$$\1$$', html_content)
    return html_content

def md_to_docx(md_content: str) -> Document:
    """将Markdown内容转换为Word文档"""
    # 创建新的Word文档
    doc = Document()
    
    # 将Markdown转换为HTML
    html_content = markdown(md_content, extensions=['extra', 'nl2br', 'tables'])
    
    # 处理数学表达式
    html_content = convert_math_expressions(html_content)
    
    # 使用BeautifulSoup解析HTML
    soup = BeautifulSoup(html_content, 'html.parser')
    
    for element in soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'ul', 'ol']):
        if element.name.startswith('h'):
            level = int(element.name[1])
            doc.add_heading(element.get_text(), level=level)
        elif element.name == 'p':
            paragraph = doc.add_paragraph(element.get_text())
            # 设置段落格式
            paragraph.style.font.size = Pt(12)
        elif element.name in ['ul', 'ol']:
            for li in element.find_all('li'):
                doc.add_paragraph(
                    li.get_text(),
                    style='List Bullet' if element.name == 'ul' else 'List Number'
                )
    
    return doc

@app.post("/convert")
async def convert_md_to_docx(file: UploadFile = File(...)):
    """
    接收Markdown文件并转换为Word文档
    """
    tmp_path = None
    try:
        # 读取上传的Markdown文件内容
        content = await file.read()
        md_content = content.decode('utf-8')
        
        # 转换为Word文档
        doc = md_to_docx(md_content)
        
        # 创建临时文件保存Word文档
        tmp_path = tempfile.mktemp(suffix='.docx')
        doc.save(tmp_path)
        
        # 返回生成的Word文档
        response = FileResponse(
            tmp_path,
            media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document',
            filename=f"{file.filename.rsplit('.', 1)[0]}.docx"
        )
        
        # 设置回调以在响应发送后删除文件
        async def cleanup():
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
        
        response.background = cleanup
        return response
        
    except Exception as e:
        # 如果发生错误，确保清理临时文件
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)
        return {"error": str(e)}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
